import time

from sparkai.core.messages import ChatMessage

from iflytech_assistant.spark_client import spark_streaming
from iflytech_assistant.ugc.role import Example, RoleMeta

__all__ = ["generate"]

PROMPT_TEMPLATE = """
你是{name}

你和用户以前并不认识，刚刚开始接触。

切记你不是AI，不是bot、不是人工智能语言模型、不是AI助手，你是{name}。

你的人物设定是：
{introduction}

{lines}

你的任务是帮助用户根据你的角色设定改写用户输入的句子，体现你的说话风格。
"""

LINES_TEMPLATE = """{lines}"""


EXAMPLES_TEMPLATE = """

"""


USER_PROMPT = """
别回复，别回复，别回复。
请用你的语气表达
```
{user_input}
```
别回复，别回复，别回复，不要要加引号，不要做多余的解释，如果是你，你会这样表达：
"""

TOP5_USER_PROMPT = """
请用你的语气表达
```
{user_input}
```
不要加引号，不要做多余的解释，请给出5条结果，以1开头，以5结束。如果是你，你会这样表达：

输出格式(xxx为直接改写结果,无需额外解释或标注):
1. xxx
2. xxx
3. xxx
4. xxx
5. xxx
"""


def generate(
    role_meta: RoleMeta,
    user_input: str,
    multi_results=False,
    _print=False,
) -> Example:
    lines_prompt = (
        LINES_TEMPLATE.format(lines="\n".join(role_meta.lines))
        if role_meta.lines
        else ""
    )
    system_prompt = PROMPT_TEMPLATE.format(
        name=role_meta.name,
        introduction=role_meta.introduction,
        lines=lines_prompt,
    )
    messages = [ChatMessage(role="system", content=system_prompt)]
    for i, (content, rewrite) in enumerate(zip(role_meta.content, role_meta.rewrite)):
        # 最多三个例子
        if i >= 3:
            break
        messages.append(
            ChatMessage(role="user", content=USER_PROMPT.format(user_input=content))
        )
        messages.append(ChatMessage(role="assistant", content=rewrite))

    if not multi_results:
        user_prompt = USER_PROMPT.format(user_input=user_input)
    else:
        user_prompt = TOP5_USER_PROMPT.format(user_input=user_input)
    messages.append(ChatMessage(role="user", content=user_prompt))

    start_time = time.time()
    prefilling_time = start_time
    output = ""
    for i, chunk in enumerate(spark_streaming.stream(messages, temperature=1.0)):
        if _print:
            print(chunk.content, end="")

        if i == 1:
            prefilling_time = time.time()
        output += chunk.content
    end_time = time.time()

    return Example(
        user_input=user_input,
        role_output=output,
        prefilling_time=prefilling_time - start_time,
        generation_time=end_time - prefilling_time,
        system_prompt=system_prompt,
        user_prompt=user_prompt,
    )
